# scripts/setup_rl_environment.py
"""
设置强化学习交易环境脚本
"""

import os
import sys
import subprocess
from pathlib import Path

def install_requirements():
    """安装依赖包"""
    print("🔧 安装Python依赖包...")
    
    requirements_file = Path("requirements_rl.txt")
    if requirements_file.exists():
        try:
            subprocess.run([
                sys.executable, "-m", "pip", "install", "-r", str(requirements_file)
            ], check=True)
            print("✅ 依赖包安装成功")
        except subprocess.CalledProcessError as e:
            print(f"❌ 依赖包安装失败: {e}")
            return False
    else:
        print("❌ requirements_rl.txt 文件不存在")
        return False
    
    return True

def create_directories():
    """创建必要的目录"""
    print("📁 创建项目目录...")
    
    directories = [
        "data",
        "models",
        "logs",
        "charts",
        "checkpoints",
        "config",
        "tests",
        "scripts"
    ]
    
    for directory in directories:
        Path(directory).mkdir(exist_ok=True)
        print(f"   ✅ {directory}/")
    
    return True

def check_gpu_availability():
    """检查GPU可用性"""
    print("🔍 检查GPU可用性...")
    
    try:
        import torch
        if torch.cuda.is_available():
            gpu_count = torch.cuda.device_count()
            gpu_name = torch.cuda.get_device_name(0)
            print(f"   ✅ 检测到 {gpu_count} 个GPU")
            print(f"   🔥 GPU型号: {gpu_name}")
            return True
        else:
            print("   ⚠️ 未检测到GPU，将使用CPU训练")
            return False
    except ImportError:
        print("   ❌ PyTorch未安装")
        return False

def download_sample_data():
    """下载示例数据"""
    print("📊 下载示例数据...")
    
    try:
        import yfinance as yf
        from datetime import datetime, timedelta
        
        # 下载示例股票数据
        symbols = ["AAPL", "GOOGL", "MSFT", "TSLA", "AMZN"]
        end_date = datetime.now()
        start_date = end_date - timedelta(days=365*2)
        
        data_dir = Path("data")
        
        for symbol in symbols:
            try:
                data = yf.download(
                    symbol,
                    start=start_date.strftime('%Y-%m-%d'),
                    end=end_date.strftime('%Y-%m-%d')
                )
                
                if not data.empty:
                    # 重命名列名
                    data.columns = [col.lower() for col in data.columns]
                    data.to_csv(data_dir / f"{symbol}_data.csv")
                    print(f"   ✅ {symbol}: {len(data)}条记录")
                else:
                    print(f"   ⚠️ {symbol}: 数据为空")
                    
            except Exception as e:
                print(f"   ❌ {symbol}: 下载失败 - {e}")
        
        print("📊 示例数据下载完成")
        return True
        
    except ImportError:
        print("   ❌ yfinance未安装，无法下载数据")
        return False
    except Exception as e:
        print(f"   ❌ 数据下载失败: {e}")
        return False

def create_config_files():
    """创建配置文件"""
    print("⚙️ 创建配置文件...")
    
    config_dir = Path("config")
    
    # 创建基本配置文件
    basic_config = """# 强化学习交易系统配置

# 环境设置
INITIAL_BALANCE = 100000
TRANSACTION_COST = 0.001
MAX_POSITION = 1.0
LOOKBACK_WINDOW = 30

# 训练设置
TRAIN_EPISODES = 1000
BATCH_SIZE = 32
LEARNING_RATE = 0.001

# 数据设置
DATA_SYMBOL = "AAPL"
TRAIN_RATIO = 0.8
"""
    
    with open(config_dir / "settings.py", "w", encoding="utf-8") as f:
        f.write(basic_config)
    
    print("   ✅ settings.py")
    
    # 创建日志配置
    log_config = """import logging
import sys
from pathlib import Path

# 创建日志目录
log_dir = Path("logs")
log_dir.mkdir(exist_ok=True)

# 配置日志
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(log_dir / "rl_trading.log"),
        logging.StreamHandler(sys.stdout)
    ]
)

# 创建logger
logger = logging.getLogger("rl_trading")
"""
    
    with open(config_dir / "logging_config.py", "w", encoding="utf-8") as f:
        f.write(log_config)
    
    print("   ✅ logging_config.py")
    
    return True

def setup_tensorboard():
    """设置TensorBoard"""
    print("📈 设置TensorBoard...")
    
    try:
        import tensorboard
        logs_dir = Path("logs/tensorboard")
        logs_dir.mkdir(parents=True, exist_ok=True)
        
        print("   ✅ TensorBoard日志目录已创建")
        print(f"   💡 启动命令: tensorboard --logdir={logs_dir}")
        return True
        
    except ImportError:
        print("   ❌ TensorBoard未安装")
        return False

def run_environment_test():
    """运行环境测试"""
    print("🧪 运行环境测试...")
    
    try:
        # 运行基本测试
        subprocess.run([
            sys.executable, "-m", "pytest", "tests/", "-v"
        ], check=True)
        print("   ✅ 所有测试通过")
        return True
        
    except subprocess.CalledProcessError:
        print("   ❌ 部分测试失败")
        return False
    except FileNotFoundError:
        print("   ⚠️ pytest未安装，跳过测试")
        return True

def main():
    """主设置函数"""
    print("🚀 强化学习交易环境设置")
    print("=" * 50)
    
    setup_steps = [
        ("安装依赖包", install_requirements),
        ("创建目录", create_directories),
        ("检查GPU", check_gpu_availability),
        ("创建配置文件", create_config_files),
        ("设置TensorBoard", setup_tensorboard),
        ("下载示例数据", download_sample_data),
        ("运行测试", run_environment_test)
    ]
    
    success_count = 0
    total_steps = len(setup_steps)
    
    for step_name, step_func in setup_steps:
        print(f"\n📍 {step_name}...")
        try:
            if step_func():
                success_count += 1
                print(f"   ✅ {step_name} 完成")
            else:
                print(f"   ⚠️ {step_name} 部分完成")
        except Exception as e:
            print(f"   ❌ {step_name} 失败: {e}")
    
    print(f"\n🎯 设置完成: {success_count}/{total_steps} 步骤成功")
    
    if success_count >= total_steps * 0.8:
        print("✅ 环境设置基本成功！")
        print("\n🎮 快速开始:")
        print("   python examples/reinforcement_learning_trading_torch.py")
        print("   python examples/reinforcement_learning_trading_torch.py --demo")
    else:
        print("⚠️ 环境设置不完整，请检查错误信息")
    
    return success_count >= total_steps * 0.6

if __name__ == "__main__":
    main()
